
from matplotlib import pyplot as plt
import warnings
import seaborn as sns

from utils.data_split import *
warnings.filterwarnings("ignore")
import statsmodels.api as sm
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score
import numpy as np
from sklearn.metrics import (
    roc_curve, auc,
    accuracy_score, f1_score, precision_score, recall_score,
    roc_auc_score, confusion_matrix
)
best_threshold_by_youden = 0.4291484596852248

warnings.filterwarnings("ignore")
def bootstrap_metric_ci(y_true, y_scores, metric_fn, threshold=None, n_bootstraps=200, ci=0.95, random_state=42):
    """
    使用 bootstrap 计算二分类指标的置信区间
    y_true: array-like, 真实标签
    y_scores: array-like, 预测概率
    metric_fn: 函数，返回单个指标（如 accuracy, sensitivity 等）
            signature: metric_fn(y_true, y_scores, threshold)
    threshold: 分类阈值
    n_bootstraps: 重采样次数
    ci: 置信水平
    """
    rng = np.random.RandomState(random_state)
    bootstrapped_scores = []
    n = len(y_true)
    
    for _ in range(n_bootstraps):
        indices = rng.choice(np.arange(n), size=n, replace=True)
        if len(np.unique(y_true[indices])) < 2:
            continue
        if threshold is not None:
            score = metric_fn(y_true[indices], y_scores[indices], threshold)
        else:
            score = metric_fn(y_true[indices], y_scores[indices])
        bootstrapped_scores.append(score)
    
    alpha = (1 - ci) / 2
    lower = np.percentile(bootstrapped_scores, 100 * alpha)
    upper = np.percentile(bootstrapped_scores, 100 * (1 - alpha))
    return lower, upper
def compute_metrics(y_true, y_scores, threshold):
    y_pred = (y_scores >= threshold).astype(int)
    acc = accuracy_score(y_true, y_pred)
    sens = recall_score(y_true, y_pred)  # TPR
    tn, fp, fn, tp = confusion_matrix(y_true, y_pred).ravel()
    spec = tn / (tn + fp) if (tn + fp) > 0 else 0
    ppv  = tp / (tp + fp) if (tp + fp) > 0 else 0
    npv  = tn / (tn + fn) if (tn + fn) > 0 else 0
    f1   = f1_score(y_true, y_pred)
    return acc, sens, spec, ppv, npv, f1
# ======================
# 📂 1. 读取数据
# ======================
signature_data = pd.read_csv(signature_score_csv_path)   # 包含 ensemble_score
ensemble_model = joblib.load(opj(ensemble_model_weight_path,'model.joblib'))
dp = DataSplitUtil(split_random_state=split_random_state_list[0])
signature_train, signature_test = dp.get_train_test_df(signature_data)
signature_test_X = signature_test.drop(columns=exclude_columns)
signature_test_y = signature_test[target_column]


signature_result = pd.DataFrame({
    'Patient_ID': signature_test['Patient_ID'].values,
    target_column: signature_test[target_column].values,
    'ensemble_score': ensemble_model.predict_proba_by_dif(signature_test_X)[:,1]  # 这里直接是 1D 数组，长度 = n_samples
})
clinic_data = pd.read_csv(opj(base_path,'data/original_data/BCNB/Clinic/patient-clinical-data-process.csv'))
df = pd.merge(signature_result, clinic_data, on=['Patient_ID','ALN status'])


# ======================
# 🧠 2. 定义子组分析函数（集成模型版）
# ======================
def bootstrap_auc_ci(y_true, y_scores, n_bootstraps=200, ci=0.95, random_state=42):
    rng = np.random.RandomState(random_state)
    bootstrapped_scores = []
    n = len(y_true)
    for _ in range(n_bootstraps):
        indices = rng.choice(np.arange(n), size=n, replace=True)
        if len(np.unique(y_true.iloc[indices])) < 2:
            continue
        fpr, tpr, _ = roc_curve(y_true.iloc[indices], y_scores.iloc[indices])
        score = auc(fpr, tpr)
        bootstrapped_scores.append(score)
    alpha = (1 - ci) / 2
    lower = np.percentile(bootstrapped_scores, 100 * alpha)
    upper = np.percentile(bootstrapped_scores, 100 * (1 - alpha))
    return lower, upper

def run_subgroup_auc_ensemble(df, score_col, group_col, group_val, label_col='ALN status'):
    """
    针对集成模型预测结果做亚组分析
    """
    # --- 选出亚组 ---
    test_sub = df[group_val(df[group_col])]
    if test_sub[label_col].nunique() < 2:
        return None

    n_alnm = sum(test_sub[label_col] == 1)
    n_non_alnm = sum(test_sub[label_col] == 0)

    # --- 计算 AUC ---
    y_true = test_sub[label_col]
    y_score = test_sub[score_col]
    y_true = y_true.reset_index(drop=True)
    y_score = y_score.reset_index(drop=True)
    auc = roc_auc_score(y_true, y_score)
    ci_low, ci_up = bootstrap_auc_ci(y_true, y_score)

    acc, sens, spec, ppv, npv, f1 = compute_metrics(y_true.values, y_score, best_threshold_by_youden)
    
    acc_ci = bootstrap_metric_ci(y_true.values, y_score,
                                lambda y, s, t: compute_metrics(y, s, t)[0],
                                threshold=best_threshold_by_youden)
    sens_ci = bootstrap_metric_ci(y_true.values, y_score,
                                lambda y, s, t: compute_metrics(y, s, t)[1],
                                threshold=best_threshold_by_youden)
    spec_ci = bootstrap_metric_ci(y_true.values, y_score,
                                lambda y, s, t: compute_metrics(y, s, t)[2],
                                threshold=best_threshold_by_youden)
    ppv_ci = bootstrap_metric_ci(y_true.values, y_score,
                                lambda y, s, t: compute_metrics(y, s, t)[3],
                                threshold=best_threshold_by_youden)
    npv_ci = bootstrap_metric_ci(y_true.values, y_score,
                                lambda y, s, t: compute_metrics(y, s, t)[4],
                                threshold=best_threshold_by_youden)
    f1_ci = bootstrap_metric_ci(y_true.values, y_score,
                                lambda y, s, t: compute_metrics(y, s, t)[5],
                                threshold=best_threshold_by_youden)
    acc_str = f"{acc:.3f} [{acc_ci[0]:.3f},{acc_ci[1]:.3f}]"
    sens_str = f"{sens:.3f} [{sens_ci[0]:.3f},{sens_ci[1]:.3f}]"
    spec_str = f"{spec:.3f} [{spec_ci[0]:.3f},{spec_ci[1]:.3f}]"
    ppv_str = f"{ppv:.3f} [{ppv_ci[0]:.3f},{ppv_ci[1]:.3f}]"
    npv_str = f"{npv:.3f} [{npv_ci[0]:.3f},{npv_ci[1]:.3f}]"
    f1_str = f"{f1:.3f} [{f1_ci[0]:.3f},{f1_ci[1]:.3f}]"
    auc_str = f"{auc:.3f} [{ci_low:.3f},{ci_up:.3f}]"

    # --- Logistic 回归 OR ---
    X_sm = sm.add_constant(y_score)
    y = y_true
    sm_model = sm.Logit(y, X_sm).fit(disp=0)

    coef = sm_model.params[1]
    conf = sm_model.conf_int().iloc[1]
    p_value = sm_model.pvalues[1]

    or_val = np.exp(coef)
    or_lower = np.exp(conf[0])
    or_upper = np.exp(conf[1])




    return {
        'Signature': 'Ensemble',
        'Group': f"{group_col}:{group_val.__name__ if hasattr(group_val,'__name__') else str(group_val)}",
        'N_test': len(test_sub),
        'AUC': auc_str,
        'OR': or_val,
        'p': p_value,
        'lower': or_lower,
        'upper': or_upper,
        'N_non_ALNM': n_non_alnm,
        'N_ALNM': n_alnm,
        'Accuracy': acc_str,
        'Sensitivity': sens_str,
        'Specificity': spec_str,
        'PPV': ppv_str,
        'NPV': npv_str,
        'F1': f1_str,
    }

# ======================
# ⚙️ 3. 定义子组条件
# ======================
subgroup_conditions = {
    'Age(years)': {
        '<=40': lambda x: x <= 40,
        '>40': lambda x: x > 40
    },
    'Tumour Size(cm)': {
        '<=2cm': lambda x: x <= 2,
        '>2cm': lambda x: x > 2
    },
    'Tumour Type': {
        'IDC': lambda x: x == 'Invasive ductal carcinoma',
        'Other': lambda x: x.isin(['Invasive lobular carcinoma', 'Other type'])
    },
    'Molecular subtype': {
        'Luminal A': lambda x: x == 'Luminal A',
        'Luminal B': lambda x: x == 'Luminal B',
        'HER2+': lambda x: x == 'HER2(+)',
        'TNBC': lambda x: x == 'Triple negative'
    }
}

# ======================
# 📊 4. 运行分析
# ======================
subgroup_results = []
for group_var, group_dict in subgroup_conditions.items():
    for name, cond in group_dict.items():
        r = run_subgroup_auc_ensemble(
            df=df,
            score_col='ensemble_score',
            group_col=group_var,
            group_val=cond,
            label_col='ALN status'
        )
        if isinstance(r, dict):
            r['Group'] = f"{group_var}:{name}"
            subgroup_results.append(r)


# 合并结果
results_df = pd.DataFrame(subgroup_results)
results_df['Subgroup_Var'] = results_df['Group'].apply(lambda x: x.split(':')[0])

# 保存结果
md(subgroup_analysis_result_path)
need_list = ['Group','AUC','Accuracy','Sensitivity','Specificity','PPV','NPV','F1','p']
results_df = results_df[need_list+ [c for c in results_df.columns if c not in need_list]]
results_df.to_csv(opj(subgroup_analysis_result_path,'ensemble_subgroup_results.csv'), index=False, encoding='utf-8-sig')

# ======================
# 📈 5. 绘制 AUC 柱状图
# ======================
# subgroup_vars = results_df['Subgroup_Var'].unique()
# for var in subgroup_vars:
#     plt.figure(figsize=(10, 6))
#     sub_df = results_df[results_df['Subgroup_Var'] == var].copy()
#     sub_df['Category'] = sub_df.apply(lambda row: f"{row['Group'].split(':')[1]} (n={int(row['N_test'])})", axis=1)
    # ax = sns.barplot(data=sub_df, x='Category', y='AUC', color="steelblue")

    # for container in ax.containers:
    #     ax.bar_label(container, fmt='%.2f', label_type='edge', fontsize=9, padding=2)

    # plt.ylim(0.2, 1.0)
    # plt.title(f"AUC of Ensemble Model by Subgroup: {var}")
    # plt.xticks(rotation=0)
    # plt.tight_layout()

    # safe_varname = var.replace(' ', '_').replace('(', '').replace(')', '')
    # save_path = opj(opj(subgroup_analysis_result_path, 'barplot', f"{safe_varname}.pdf"))
    # md(opd(save_path))
    # plt.savefig(save_path, dpi=300)